import os
import pickle
import torch
import numpy as np

from tqdm import trange
from PIL import Image
from segbase import SegmentationDataset


class CoronarySegmentation(SegmentationDataset):
    """coronary_seg_data Dataset
    """
    NUM_CLASS = 2

    def __init__(self, root='/home/handewei/project/segment/awesome-segmentation/datasets', split='train', mode=None, transform=None, **kwargs):
        super(CoronarySegmentation, self).__init__(root, split, mode, transform, **kwargs)

        with open(os.path.join(root, 'coronary.txt'), 'r') as f:
            self.image_all_list = f.readlines()
        one_tenth = len(self.image_all_list) // 10
        if self.mode == 'train':
            self.img_list = self.image_all_list[0: one_tenth * 7]
        elif self.mode == 'val':
            self.img_list = self.image_all_list[one_tenth * 7: one_tenth * 8]
        elif self.mode == 'test':
            self.img_list = self.image_all_list[one_tenth * 8: len(self.image_all_list)]

    def __getitem__(self, index):
        image = Image.open(self.img_list[index].replace('\n', '')).convert('RGB')
        label = Image.open(self.img_list[index].replace("src", "label").replace('\n', ''))

        if self.mode == 'test':
            if self.transform is not None:
                image = self.transform(image)
            return image, os.path.basename(self.img_list[index])

        # synchrosized transform
        if self.mode == 'train':
            image, label = self._sync_transform(image, label)
        elif self.mode == 'val':
            image, label = self._val_sync_transform(image, label)
        elif self.mode == 'test':
            image, label = self._img_transform(image), self._mask_transform(label)

        # general resize, normalize and toTensor
        if self.transform is not None:
            image = self.transform(image)
        return image, label, os.path.basename(self.img_list[index])

    def _mask_transform(self, mask):
        target = np.array(mask).astype('int32')
        target[target > 0] = 1
        return torch.from_numpy(target).long()

    def __len__(self):
        return len(self.img_list)

    @property
    def pred_offset(self):
        return 0

if __name__ == '__main__':
    dataset = CoronarySegmentation(base_size=520, crop_size=520)